/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: offline build cluster info parse
 * Author: wangke
 * Create: 2022-10-09
 */

#include <nlohmann/json.hpp>
#include "offline_build_config_parse.h"
#include <securec.h>
#include <functional>
#include <vector>
#include <memory>
#include <mutex>

#include "graph/utils/tensor_utils.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/utils/node_utils.h"
#include "graph/utils/graph_utils.h"
#include "graph/utils/op_desc_utils.h"
#include "framework/memory/memory_api.h"
#include "graph/ge_local_context.h"
#include "framework/common/ge_types.h" // ge对外options
#include "hccl/hcom.h"
#include "externalinput_pub.h"
#include "hcom_pub.h"
#include "hcom_op_utils.h"
#include "ops_kernel_builder_base.h"
#include "op_hcom_comm.h"
#include "param_check_pub.h"
#include "config.h"
#include "workflow_pub.h"
#include "sal.h"
#include "coll_alg_utils.h"
#include "adapter_rts.h"

using namespace std;

static std::unordered_map<s32, u64> OFFLINE_BUILD_SUB_STEAM_NUM = {
    {HCCL_DEVICE_NUM_EIGHT, HCCL_SUB_STREAM_NUM_8P_RING},
    {HCCL_DEVICE_NUM_FOUR, HCCL_SUB_STREAM_NUM_4P_MESH},
    {HCCL_DEVICE_NUM_TWO, HCCL_SUB_STREAM_NUM_ZERO},
    {HCCL_DEVICE_MINNUM, HCCL_SUB_STREAM_NUM_ZERO},
};

namespace hccl {
static std::mutex g_taskNumCalModeMutex;

bool IsOfflineCompilation()
{
    std::string offlineString;
    if (ge::GetThreadLocalContext().GetOption("ge.offline_hccl_compile", offlineString) == ge::GRAPH_SUCCESS) {
        return true;
    }
    return false;
}

HcclResult GetDeterministic(bool &isDeterministic)
{
    isDeterministic = false; // 默认为false
    DevType devType;
    CHK_RET(GetOffDeviceTypeWithoutDev(devType));
    // 910_93时，不解析确定性计算环境变量，设置为默认值
    if (devType == DevType::DEV_TYPE_910_93) {
        HCCL_WARNING("[GetDeterministic] device type is 910_93, use default setting");
        return HCCL_SUCCESS;
    }

    std::string hcclDeterministicEnv = SalGetEnv("HCCL_DETERMINISTIC");
    if (hcclDeterministicEnv != "EmptyString") {
        // 环境变量优先
        std::transform(
            hcclDeterministicEnv.begin(), hcclDeterministicEnv.end(), hcclDeterministicEnv.begin(), ::toupper);
        if ("FALSE" == hcclDeterministicEnv) {
            isDeterministic = false;
        } else if("TRUE" == hcclDeterministicEnv) {
            isDeterministic = true;
        } else {
            HCCL_ERROR("[GetDeterministic] HCCL_DETERMINISTIC is set to [%s], which is incorrect. Please check", hcclDeterministicEnv.c_str());
            return HCCL_E_PARA;
        }
    } else {
        // 未配环境变量，检查ge option
        std::string geOption;
        if (ge::GetThreadLocalContext().GetOption(ge::DETERMINISTIC, geOption) == ge::GRAPH_SUCCESS) {
            if (geOption == "1") {
                isDeterministic = true;
            }
        }
    }

    return HCCL_SUCCESS;
}

HcclResult GetAllReduceScratchSizeWithoutDev(const string &sGroup, const ge::Node& node,
    s32 serverNum, s32 rankSize, const u64 count, const HcclDataType dataType, u64 &scratchSize)
{
    // 检查是否使能确定性计算
    bool isDeterministic = false;
    CHK_RET(GetDeterministic(isDeterministic));

    std::string reduction;
    if (ge::AttrUtils::GetStr(node.GetOpDesc(), ge::HCOM_ATTR_REDUCE_TYPE, reduction) != ge::GRAPH_SUCCESS) {
        // 部分场景可能无法取得reduce类型，保守地，按非prod处理，尝试以宽松的条件去分配scratch
        HCCL_WARNING("[GetAllReduceScratchMemSize] Get ReduceType failed.");
    }
    bool supportInlineReduce = (reduction != "prod") && (dataType != HcclDataType::HCCL_DATA_TYPE_INT64);
    DevType devType;
    CHK_RET(GetOffDeviceTypeWithoutDev(devType));
    bool isAscend910B = (devType == DevType::DEV_TYPE_910B);
    u64 memSize = SIZE_TABLE[dataType] * count;
    if (isDeterministic && serverNum <= 1 && isAscend910B && supportInlineReduce) {
        const u32 deviceEight = 8;
        const s32 deviceTwo = 2;

        scratchSize = 0;
        if (serverNum == 0) {
            // 无效serverNum，按最大需求申请
            if (memSize <= HCCL_SMALL_COUNT_GRAPH_64_KB) {
                // 小数据
                scratchSize = memSize * (deviceEight - 1);
            }
        } else if (serverNum == 1 && rankSize > deviceTwo) {
            // 有效serverNum，按实际需求申请
            if (memSize <= HCCL_SMALL_COUNT_GRAPH_64_KB) {
                // 小数据
                if (rankSize == deviceEight) {
                    scratchSize = 0; // Small Count HD
                } else {
                    scratchSize = memSize * (rankSize - 1); // Small Count Reduce+Bcast
                }
            }
        }
    }
    HCCL_DEBUG("[GetAllReduceScratchMemSizeWithoutDev] serverNum[%d], memSize[%llu], rankSize[%d], scratchSize[%llu]",
               serverNum, memSize, rankSize, scratchSize);
    return HCCL_SUCCESS;
}

HcclResult GetOffDeviceTypeWithoutDev(DevType &devType)
{
    // 离线编译第一阶段获取devType从SOC_VERSION里获取
    std::string socVersion;
    if (ge::GetThreadLocalContext().GetOption(ge::SOC_VERSION, socVersion) != ge::GRAPH_SUCCESS) {
        HCCL_ERROR("[offline][compilation] get soc version failed.");
        return HCCL_E_NOT_FOUND;
    }

    DevType tempDevType = DevType::DEV_TYPE_COUNT;
    CHK_RET(hrtGetDeviceTypeBySocVersion(socVersion, tempDevType));

    if (tempDevType != DevType::DEV_TYPE_910 && tempDevType != DevType::DEV_TYPE_910B &&
        tempDevType != DevType::DEV_TYPE_310P1 && tempDevType != DevType::DEV_TYPE_310P3 &&
        tempDevType != DevType::DEV_TYPE_910_93) {
        HCCL_ERROR("[offline][compilation] cur dev type[%u] is not support.", tempDevType);
        return HCCL_E_RUNTIME;
    }
    devType = tempDevType;
    HCCL_DEBUG("[offline] Get devtype[%u]....", devType);
    return HCCL_SUCCESS;
}

HcclResult GetSubStreamNum(s32 deviceNum, u64 &streamNum, s32 &serverNum)
{
    DevType deviceType;
    // 图编译/图优化阶段不能使用rt接口获取devicetype
    CHK_RET(GetOffDeviceTypeWithoutDev(deviceType)); // 80和71要分开
    if (deviceType == DevType::DEV_TYPE_910B) {
        constexpr u64 maxStream = 6;
        streamNum = std::min(maxStream, static_cast<u64>(deviceNum) - MINUS_MESH_STREAM_NUM);
        if (GetExternalInputHcclAlgoConfig()[HCCL_ALGO_LEVEL_1] == HcclAlgoType::HCCL_ALGO_TYPE_PIPELINE) {
            streamNum = static_cast<u64>(deviceNum);
        }
    } else if (deviceType == DevType::DEV_TYPE_910_93) {
        if (serverNum == 1) {
            streamNum = static_cast<u64>(deviceNum) - MINUS_MESH_STREAM_NUM;
        } else {
            constexpr u64 streamNumFor91093 = 3;
            streamNum = streamNumFor91093;
        }
    } else {
        if (deviceNum > HCCL_DEVICE_NUM_EIGHT) {
            streamNum = OFFLINE_BUILD_SUB_STEAM_NUM[HCCL_DEVICE_NUM_EIGHT];
        } else if (OFFLINE_BUILD_SUB_STEAM_NUM.count(deviceNum) != 0) {
            streamNum = OFFLINE_BUILD_SUB_STEAM_NUM[deviceNum];
        } else {
            streamNum = 0;
        }
    }

    if (SatisfyIntraSuperPod(deviceType, deviceNum, true)) {
        streamNum = std::max(static_cast<u64>(deviceNum - 1u), streamNum);
    } else if (FullmeshPairwiseSatisfyHighPerfAlltoallMeshCondition(deviceType, deviceNum * serverNum, true)) {
        streamNum = std::max(static_cast<u64>(deviceNum * serverNum - 1u), streamNum);
    }

    HCCL_DEBUG("[GetSubStreamNum]deviceNum[%d], streamNum[%llu]", deviceNum, streamNum);
    return HCCL_SUCCESS;
}

HcclResult GetServerIdByRankId(const std::vector<RankInfo_t> &rankList, const u32 &rankId, u32 &serverId)
{
    for (auto &iter : rankList) {
        if (iter.rankId == rankId) {
            serverId = iter.serverIdx;
            return HCCL_SUCCESS;
    }
    }
    return HCCL_E_NOT_FOUND;
}

HcclResult GetServerAndDevNumFromRanklist(const std::vector<u32> &groupRanks, const std::vector<RankInfo_t> &rankList,
    s32 &serverNum, s32 &deviceNum, bool &multiModuleDiffDeviceNumMode)
{
    u32 serverId = 0;
    std::map<u32, s32> serverAndDevNum;
    deviceNum = 0;
    for (auto it : groupRanks) {
        CHK_RET(GetServerIdByRankId(rankList, it, serverId));
        if (serverAndDevNum.find(serverId) == serverAndDevNum.end()) {
            serverAndDevNum[serverId] = 1;
        } else {
            serverAndDevNum[serverId]++;
        }
        // 可能存在不同server内device数量不一致的情况，因此求最大值
        if (serverAndDevNum[serverId] > deviceNum) {
            deviceNum = serverAndDevNum[serverId];
        }
    }
    serverNum = serverAndDevNum.size();

    // 获取multiModuleDiffDeviceNumMode信息
    HcclResult ret = HcomOpUtils::GetModuleInfo(rankList, multiModuleDiffDeviceNumMode);
    if (ret != HCCL_SUCCESS) {
        HCCL_WARNING("call GetModuleInfo error, failed to get multiModuleDiffDeviceNumMode.");
    }
    return HCCL_SUCCESS;
}

HcclResult GetServerAndDevNumFromLogRanktable(const string &ranktableString, const std::vector<u32> &groupRanks,
    s32 &serverNum, s32 &deviceNum, bool &multiModuleDiffDeviceNumMode)
{
    HCCL_INFO("Entry GetServerAndDevNumFromLogRanktable");
    RankTable_t clusterInfo;
    TopoinfoRanktableOffline myTopoRanktable(ranktableString);
    CHK_RET(myTopoRanktable.Init());
    CHK_RET(myTopoRanktable.GetClusterInfo(clusterInfo));

    CHK_RET(GetServerAndDevNumFromRanklist(groupRanks, clusterInfo.rankList, serverNum, deviceNum,
        multiModuleDiffDeviceNumMode));
    return HCCL_SUCCESS;
}

HcclResult GetServerAndDevNumFromRankMapping(const string &rankmappingString, const std::vector<u32> &groupRanks,
    s32 &serverNum, s32 &deviceNum)
{
    HCCL_INFO("Entry GetServerAndDevNumFromRankMapping");
    nlohmann::json mappingList;
    CHK_RET(SalParseInformation(mappingList, rankmappingString));
    std::map<u32, s32> serverAndDevNum;
    deviceNum = 0;
    for (auto it : groupRanks) {
        for (auto &mappingInfo : mappingList) {
            if (mappingInfo["rank_id"] == std::to_string(it)) {
                vector<u32> curMappings = mappingInfo["device_index"];
                if (serverAndDevNum.find(curMappings[1]) == serverAndDevNum.end()) {
                    serverAndDevNum[curMappings[1]] = 1;
                } else {
                    serverAndDevNum[curMappings[1]]++;
                }
                if (serverAndDevNum[curMappings[1]] > deviceNum) {
                    deviceNum = serverAndDevNum[curMappings[1]];
                }
            }
        }
    }
    serverNum = serverAndDevNum.size();
    return HCCL_SUCCESS;
}

HcclResult GetServerAndDevNumFromPhyRanktable(const string &ranktableString, const std::vector<u32> &groupRanks,
    s32 &serverNum, s32 &deviceNum, bool &multiModuleDiffDeviceNumMode)
{
    HCCL_INFO("Entry GetServerAndDevNumFromPhyRanktable");
    HcclResult ret = HCCL_SUCCESS;
    HcomInfo hcomInfo;
    std::string rankTableM;
    std::string realFilePath;
    std::string identify = "0";
    u32 rankTableSize = 0;
    ret = HcomLoadRanktableFile(ranktableString.c_str(), rankTableM, realFilePath);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[GetServerAndDevNumFromPhyRanktable]errNo[0x%016llx] rankTablePath[%s]"
        "load rankTable error.", HCCL_ERROR_CODE(HCCL_E_INTERNAL), ranktableString.c_str()), HCCL_E_INTERNAL);
    ret = HcomCheckRankTable(rankTableM.c_str(), rankTableSize);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Init][Result]errNo[0x%016llx] input rankTable error", HCOM_ERROR_CODE(ret)), ret);
    std::unique_lock<std::mutex> lock(g_taskNumCalModeMutex);
    SetTaskNumCalMode(true);
    ret = CfgGetClusterInfo(rankTableM, identify, hcomInfo.params, hcomInfo.rankTable);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Init][Result]errNo[0x%016llx] cfg get ranktable[%s] info "\
        "error: identify[%s]", HCOM_ERROR_CODE(ret), rankTableM.c_str(), identify.c_str()), ret);
    SetTaskNumCalMode(false);
    lock.unlock();

    CHK_RET(GetServerAndDevNumFromRanklist(groupRanks, hcomInfo.rankTable.rankList, serverNum, deviceNum,
        multiModuleDiffDeviceNumMode));
    return HCCL_SUCCESS;
}

HcclResult GetServerAndDevNumFromGroup(const std::vector<u32> &groupRanks, s32 &serverNum, s32 &deviceNum,
    bool &multiModuleDiffDeviceNumMode)
{
    string optionString;
    if ((ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_HCOM_RANK_MAPPING, optionString)
        == ge::GRAPH_SUCCESS) && !optionString.empty()) {
        CHK_RET(GetServerAndDevNumFromRankMapping(optionString, groupRanks, serverNum, deviceNum));
    } else if ((ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_RANK_TABLE, optionString) ==
        ge::GRAPH_SUCCESS) && !optionString.empty()) {
        CHK_RET(GetServerAndDevNumFromLogRanktable(optionString, groupRanks, serverNum, deviceNum,
            multiModuleDiffDeviceNumMode));
    } else if ((ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_RANK_TABLE_FILE, optionString) ==
        ge::GRAPH_SUCCESS) && !optionString.empty()) {
        CHK_RET(GetServerAndDevNumFromPhyRanktable(optionString, groupRanks, serverNum, deviceNum,
            multiModuleDiffDeviceNumMode));
    } else {
        HCCL_ERROR("logic ranktable and rank mapping and physical ranktable is not exist,\
            please Check the input parameters.");
        return HCCL_E_PARA;
    }
    return HCCL_SUCCESS;
}

HcclResult GetDeviceAndServerNumFromGroupList(const std::string &groupListString, const string groupName,
    s32 &serverNum, s32 &deviceNumPerServer, bool &multiModuleDiffDeviceNumMode)
{
    HCCL_INFO("get device num from group list, group name is [%s]", groupName.c_str());
    if (groupName == HCCL_WORLD_GROUP) {
        deviceNumPerServer = 0;
        serverNum = 0;
        return HCCL_SUCCESS;
    }

    try {
        // 获取并设定stream 数量
        // 能获取到group list时进入离线编译的流程去获取从流个数
        deviceNumPerServer = 0;
        nlohmann::json groupListConf;
        CHK_RET(SalParseInformation(groupListConf, groupListString));
        std::vector<nlohmann::json> groupList = groupListConf.get<std::vector<nlohmann::json>>();
        for (auto &groupInfo : groupList) {
            HCCL_DEBUG("groupInfo:%s", groupInfo.dump().c_str());
            std::string curGroupName = groupInfo["group_name"];
            HCCL_DEBUG("curGroupName:%s", curGroupName.c_str());
            if (curGroupName == groupName) {
                std::vector<u32> curRanks = groupInfo["group_rank_list"];
                CHK_RET(GetServerAndDevNumFromGroup(curRanks, serverNum, deviceNumPerServer, multiModuleDiffDeviceNumMode));
                HCCL_INFO("deviceNumPerServer:[%d] serverNum:[%d]", deviceNumPerServer, serverNum);
                break;
            }
        }
    } catch (const std::exception& e) {
        HCCL_ERROR("[HcomCalcOpRunningParam] exception caught. err[%s]", e.what());
        return HCCL_E_INTERNAL;
    }

    CHK_PRT_RET(deviceNumPerServer == 0,
        HCCL_ERROR("[HcomCalcOpRunningParam]cur group size is 0, maybe groupList is incomplete"), HCCL_E_PARA);
    HCCL_INFO("get device num per server[%d] server num[%d] from grouplist", deviceNumPerServer, serverNum);
    return HCCL_SUCCESS;
}

HcclResult GetClusterInfoAndDeviceNum(RankTable_t &clusterInfo, s32 &deviceNum)
{
    std::string rankMappingString;

    if (ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_RANK_TABLE, rankMappingString) !=
        ge::GRAPH_SUCCESS) {
        HCCL_ERROR("[get][offlineStreamNum] get ge::OPTION_EXEC_RANK_TABLE failed.");
        return HCCL_E_NOT_FOUND;
    }
    HCCL_DEBUG("[get][offlineStreamNum]rankMappingString[%s]", rankMappingString.c_str());
    TopoinfoRanktableOffline myTopoRanktable(rankMappingString);
    CHK_RET(myTopoRanktable.Init());
    CHK_RET(myTopoRanktable.GetClusterInfo(clusterInfo));
    CHK_RET(myTopoRanktable.GetDeviceNumPerServer(deviceNum));
    CHK_PRT_RET(deviceNum == 0, HCCL_ERROR("[GetStremNumOfflineByDev]cur device num per server is 0,\
        maybe ranktable is incomplete"), HCCL_E_PARA);
    return HCCL_SUCCESS;
}
HcclResult GetStremNumOfflineByDev(const DevType &devType, const std::string &sCollectiveType,
    const std::string curGroup, u64 &streamNum)
{
    // 先根据逻辑ranktable获取device数量
    s32 deviceNumPerServer = 0;
    s32 serverNum = 0;
    bool multiModuleDiffDeviceNumMode = false;
    RankTable_t clusterInfo;

    if (devType == DevType::DEV_TYPE_910 && (sCollectiveType == HCCL_KERNEL_OP_TYPE_SEND ||
        sCollectiveType == HCCL_KERNEL_OP_TYPE_RECEIVE)) {
        streamNum = 0;
        return HCCL_SUCCESS;
    }

    // world group 从逻辑ranktable里获取device数量
    if (curGroup == HCCL_WORLD_GROUP) {
        CHK_RET(GetClusterInfoAndDeviceNum(clusterInfo, deviceNumPerServer));
    } else {
        std::string groupListString;
        if (ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_HCOM_GROUPLIST, groupListString) !=
            ge::GRAPH_SUCCESS) {
            HCCL_ERROR("offline compile need have OPTION_EXEC_HCOM_GROUPLIST.please make true");
            return HCCL_E_PARA;
        }
        CHK_RET(GetDeviceAndServerNumFromGroupList(groupListString, curGroup, serverNum, deviceNumPerServer,
            multiModuleDiffDeviceNumMode));
    }

    if (devType == DevType::DEV_TYPE_910 && deviceNumPerServer == HCCL_DEVICE_NUM_EIGHT) {
        CHK_RET(GetSubStreamNum(deviceNumPerServer, streamNum, serverNum));
    } else if (devType == DevType::DEV_TYPE_910_93) {
        CHK_RET(GetSubStreamNum(deviceNumPerServer, streamNum, serverNum));
    } else {
        streamNum = deviceNumPerServer > HCCL_DEVICE_NUM_ONE ? deviceNumPerServer - MINUS_MESH_STREAM_NUM : 0;
    }
    HCCL_INFO("[GetStremNumOfflineByDev] get device num per server is [%u] streamNum [%u] curGroup[%s]",
        deviceNumPerServer, streamNum, curGroup.c_str());
    return HCCL_SUCCESS;
}

HcclResult GetStreamNumOfflineComp(const std::string &sCollectiveType, const std::string curGroup, u64 &streamNum)
{
    DevType devType;

    // 获取芯片类型
    CHK_RET(GetOffDeviceTypeWithoutDev(devType));
    HCCL_INFO("[GetStreamNumOfflineComp] get dev type[%u] op type[%s]", devType, sCollectiveType.c_str());

    switch (devType) {
        case DevType::DEV_TYPE_310P1:
        case DevType::DEV_TYPE_310P3: {
            streamNum = 0;
            break;
        }

        case DevType::DEV_TYPE_910B:
        case DevType::DEV_TYPE_910: {
            CHK_RET(GetStremNumOfflineByDev(devType, sCollectiveType, curGroup, streamNum));
            break;
        }

        default: {
            HCCL_ERROR("[Get][OfflineCompStreamNum] The cur devtype[%u] does not support offline compilation",
                devType);
            return HCCL_E_NOT_SUPPORT;
        }
    }
    HCCL_INFO("[GetStreamNumOfflineComp]stream num is [%llu]", streamNum);
    return HCCL_SUCCESS;
}
}
